# services/data_service.py
import pandas as pd
import time
import logging
from datetime import datetime, timedelta
from typing import List, Optional

from stock_data_system.database.postgres_manager import PostgresManager
from stock_data_system.data_sources.baostock_client import BaostockClient
from stock_data_system.config.database import DatabaseConfig, BaostockConfig

class StockDataService:
    """股票数据服务"""

    def __init__(self, db_config: DatabaseConfig, bs_config: BaostockConfig):
        self.db_manager = PostgresManager(db_config)
        self.bs_client = BaostockClient(bs_config)
        self.logger = logging.getLogger(__name__)

    def initialize_database(self):
        """初始化数据库表结构"""
        self.db_manager.create_tables()
        self.logger.info("数据库初始化完成")

    def test_baostock_connection(self) -> bool:
        """测试Baostock连接"""
        return self.bs_client.test_connection()

    def update_stock_basic(self) -> int:
        """更新股票基本信息"""
        try:
            if not self.test_baostock_connection():
                self.logger.error("Baostock连接测试失败")
                return 0

            self.logger.info("正在从Baostock获取股票基本信息...")

            df = self.bs_client.get_stock_basic()
            if df.empty:
                self.logger.warning("未获取到股票信息")
                return 0

            # 保存到数据库
            self.db_manager.bulk_insert_dataframe(df, 'stock_basic', 'replace')

            # 显示统计信息
            market_counts = df['market'].value_counts()
            self.logger.info("股票分布统计:")
            for market, count in market_counts.items():
                self.logger.info(f"  {market}: {count} 只股票")

            self._log_update('basic', None, len(df), 'success')
            self.logger.info(f"成功更新 {len(df)} 只股票基本信息")
            return len(df)

        except Exception as e:
            self.logger.error(f"更新股票基本信息失败: {e}")
            self._log_update('basic', None, 0, 'failed', str(e))
            return 0

    def _get_trading_days_ago(self, trading_days: int = 100) -> str:
        """获取指定交易日之前的日期（近似计算）"""
        # 交易日大约占自然日的70%（5个交易日/7天）
        # 100个交易日 ≈ 100 / 5 * 7 = 140个自然日
        natural_days = int(trading_days * 1.4)  # 增加一些缓冲
        target_date = datetime.now() - timedelta(days=natural_days)
        return target_date.strftime('%Y-%m-%d')

    def update_daily_data(self, symbols: List[str], start_date: str = None, end_date: str = None) -> int:
        """更新日线数据"""
        # 设置默认日期范围：100个交易日之前到现在
        if end_date is None:
            end_date = datetime.now().strftime('%Y-%m-%d')
        if start_date is None:
            start_date = self._get_trading_days_ago(100)

        self.logger.info(f"更新日期范围: {start_date} 到 {end_date}")

        total_updated = 0
        successful_symbols = []
        failed_symbols = []

        for i, symbol in enumerate(symbols, 1):
            try:
                self.logger.info(f"处理股票 {i}/{len(symbols)}: {symbol}")

                # 获取该股票的最后更新日期
                last_date = self.db_manager.get_last_trade_date(symbol)
                actual_start_date = last_date or start_date

                # 检查是否需要更新
                if pd.to_datetime(actual_start_date) > pd.to_datetime(end_date):
                    self.logger.info(f"{symbol} 数据已是最新，跳过")
                    continue

                df = self.bs_client.get_daily_data(symbol, actual_start_date, end_date)
                if not df.empty:
                    self.db_manager.bulk_insert_dataframe(df, 'daily_quotes')
                    total_updated += len(df)
                    successful_symbols.append(symbol)
                    self.logger.info(f"更新 {symbol} 日线数据 {len(df)} 条")
                else:
                    self.logger.info(f"{symbol} 没有新数据")

            except Exception as e:
                self.logger.error(f"更新 {symbol} 日线数据失败: {e}")
                failed_symbols.append(symbol)

        # 记录统计信息
        self.logger.info(f"日线数据更新完成: 成功 {len(successful_symbols)} 只, 失败 {len(failed_symbols)} 只")
        if failed_symbols:
            self.logger.warning(f"失败的股票: {failed_symbols}")

        self._log_update('daily', None, total_updated, 'success')
        return total_updated

    def update_all_stocks_daily(self, trading_days: int = 100) -> int:
        """更新所有股票的日线数据
        Args:
            trading_days: 要更新的交易天数，默认100个交易日
        """
        try:
            # 获取所有股票代码
            query = "SELECT symbol FROM stock_basic WHERE status = '1'"
            symbols_df = self.db_manager.execute_query(query)

            if symbols_df.empty:
                self.logger.warning("没有找到有效的股票代码")
                return 0

            symbols = symbols_df['symbol'].tolist()
            end_date = datetime.now().strftime('%Y-%m-%d')
            start_date = self._get_trading_days_ago(trading_days)

            self.logger.info(f"更新所有{len(symbols)}只股票的{trading_days}个交易日数据")
            return self.update_daily_data(symbols, start_date, end_date)

        except Exception as e:
            self.logger.error(f"更新所有股票数据失败: {e}")
            return 0

    def update_recent_daily_data(self, days: int = 1) -> int:
        """更新最近几天的日线数据（快捷方法）"""
        """更新最近的自然日数据（主要用于每日更新）"""
        try:
            query = "SELECT symbol FROM stock_basic WHERE status = '1'"
            symbols_df = self.db_manager.execute_query(query)

            if symbols_df.empty:
                self.logger.warning("没有找到有效的股票代码")
                return 0

            symbols = symbols_df['symbol'].tolist()
            end_date = datetime.now().strftime('%Y-%m-%d')
            start_date = (datetime.now() - timedelta(days=days)).strftime('%Y-%m-%d')

            self.logger.info(f"更新最近{days}天数据")
            return self.update_daily_data(symbols, start_date, end_date)

        except Exception as e:
            self.logger.error(f"更新最近数据失败: {e}")
            return 0

    def get_daily_data(self, symbol: str, start_date: str, end_date: str) -> pd.DataFrame:
        """获取日线数据"""
        query = """
            SELECT * FROM daily_quotes 
            WHERE symbol = :symbol AND trade_date BETWEEN :start_date AND :end_date
            ORDER BY trade_date
        """
        params = {'symbol': symbol, 'start_date': start_date, 'end_date': end_date}
        return self.db_manager.execute_query(query, params)

    def get_stock_list(self, market: str = None) -> pd.DataFrame:
        """获取股票列表"""
        if market:
            query = "SELECT * FROM stock_basic WHERE market = :market AND status = '1'"
            params = {'market': market}
        else:
            query = "SELECT * FROM stock_basic WHERE status = '1'"
            params = {}

        return self.db_manager.execute_query(query, params)

    def _log_update(self, data_type: str, symbol: str, records: int, status: str, error_msg: str = None):
        """记录更新日志"""
        query = """
            INSERT INTO update_log 
            (data_type, symbol, records_updated, status, error_message)
            VALUES (:data_type, :symbol, :records, :status, :error_msg)
        """
        params = {
            'data_type': data_type,
            'symbol': symbol,
            'records': records,
            'status': status,
            'error_msg': error_msg
        }
        self.db_manager.execute_query(query, params)

    def update_financial_data(self, symbols: List[str] = None, years: List[int] = None, quarters: List[int] = None) -> int:
        """更新财务数据"""
        try:
            if not self.test_baostock_connection():
                self.logger.error("Baostock连接测试失败")
                return 0

            # 如果没有指定股票，获取所有上市股票
            if symbols is None:
                symbols_df = self.get_stock_list()
                if symbols_df.empty:
                    self.logger.warning("没有找到有效的股票代码")
                    return 0
                symbols = symbols_df['symbol'].tolist()

            # 设置默认的时间范围（最近3年）
            if years is None:
                current_year = datetime.now().year
                years = list(range(current_year - 2, current_year + 1))

            if quarters is None:
                quarters = [1, 2, 3, 4]  # 所有季度

            self.logger.info(f"开始更新财务数据: {len(symbols)}只股票, {len(years)}年, {len(quarters)}个季度")

            total_updated = 0
            successful_symbols = []
            failed_symbols = []

            for i, symbol in enumerate(symbols, 1):
                try:
                    self.logger.info(f"处理财务数据 {i}/{len(symbols)}: {symbol}")

                    symbol_updated = 0
                    for year in years:
                        for quarter in quarters:
                            df = self.bs_client.get_financial_data(symbol, year, quarter)
                            if not df.empty:
                                self.db_manager.bulk_insert_dataframe(df, 'financial_data')
                                symbol_updated += len(df)

                    if symbol_updated > 0:
                        total_updated += symbol_updated
                        successful_symbols.append(symbol)
                        self.logger.info(f"更新 {symbol} 财务数据 {symbol_updated} 条")
                    else:
                        self.logger.info(f"{symbol} 没有找到财务数据")

                except Exception as e:
                    self.logger.error(f"更新 {symbol} 财务数据失败: {e}")
                    failed_symbols.append(symbol)

                    # 避免请求过于频繁
                    time.sleep(0.5)

            # 记录统计信息
            self.logger.info(f"财务数据更新完成: 成功 {len(successful_symbols)} 只, 失败 {len(failed_symbols)} 只")
            if failed_symbols:
                self.logger.warning(f"失败的股票: {failed_symbols}")

            self._log_update('financial', None, total_updated, 'success')
            return total_updated

        except Exception as e:
            self.logger.error(f"更新财务数据失败: {e}")
            self._log_update('financial', None, 0, 'failed', str(e))
            return 0

    def update_dividend_data(self, symbols: List[str] = None, years: List[int] = None) -> int:
        """更新分红数据"""
        try:
            if not self.test_baostock_connection():
                self.logger.error("Baostock连接测试失败")
                return 0

            # 如果没有指定股票，获取所有上市股票
            if symbols is None:
                symbols_df = self.get_stock_list()
                if symbols_df.empty:
                    self.logger.warning("没有找到有效的股票代码")
                    return 0
                symbols = symbols_df['symbol'].tolist()

            # 设置默认的时间范围（最近5年）
            if years is None:
                current_year = datetime.now().year
                years = list(range(current_year - 4, current_year + 1))

            self.logger.info(f"开始更新分红数据: {len(symbols)}只股票, {len(years)}年")

            total_updated = 0
            successful_symbols = []
            failed_symbols = []

            for i, symbol in enumerate(symbols, 1):
                try:
                    self.logger.info(f"处理分红数据 {i}/{len(symbols)}: {symbol}")

                    df = self.bs_client.get_dividend_data(symbol, years[0] if len(years) == 1 else None)
                    if not df.empty:
                        self.db_manager.bulk_insert_dataframe(df, 'dividend_data')
                        total_updated += len(df)
                        successful_symbols.append(symbol)
                        self.logger.info(f"更新 {symbol} 分红数据 {len(df)} 条")
                    else:
                        self.logger.info(f"{symbol} 没有找到分红数据")

                except Exception as e:
                    self.logger.error(f"更新 {symbol} 分红数据失败: {e}")
                    failed_symbols.append(symbol)

                # 控制请求频率
                time.sleep(0.3)

            # 记录统计信息
            self.logger.info(f"分红数据更新完成: 成功 {len(successful_symbols)} 只, 失败 {len(failed_symbols)} 只")
            if failed_symbols:
                self.logger.warning(f"失败的股票: {failed_symbols}")

            self._log_update('dividend', None, total_updated, 'success')
            return total_updated

        except Exception as e:
            self.logger.error(f"更新分红数据失败: {e}")
            self._log_update('dividend', None, 0, 'failed', str(e))
            return 0

    def update_all_financial_data(self, years: List[int] = None, quarters: List[int] = None):
        """更新所有股票的财务数据（批量处理）"""
        symbols_df = self.get_stock_list()
        if symbols_df.empty:
            self.logger.warning("没有找到有效的股票代码")
            return 0

        symbols = symbols_df['symbol'].tolist()

        # 分批处理，避免内存溢出
        batch_size = 50
        total_updated = 0

        for i in range(0, len(symbols), batch_size):
            batch_symbols = symbols[i:i + batch_size]
            self.logger.info(f"处理批次 {i//batch_size + 1}/{(len(symbols)-1)//batch_size + 1}")

            updated = self.update_financial_data(batch_symbols, years, quarters)
            total_updated += updated

            # 批次之间休息一下
            time.sleep(1)

        return total_updated

    def update_all_dividend_data(self, years: List[int] = None):
        """更新所有股票的分红数据（批量处理）"""
        symbols_df = self.get_stock_list()
        if symbols_df.empty:
            self.logger.warning("没有找到有效的股票代码")
            return 0

        symbols = symbols_df['symbol'].tolist()

        # 分批处理
        batch_size = 50
        total_updated = 0

        for i in range(0, len(symbols), batch_size):
            batch_symbols = symbols[i:i + batch_size]
            self.logger.info(f"处理批次 {i//batch_size + 1}/{(len(symbols)-1)//batch_size + 1}")

            updated = self.update_dividend_data(batch_symbols, years)
            total_updated += updated

            time.sleep(1)

        return total_updated

    # 现有的查询方法
    def get_daily_data(self, symbol: str, start_date: str, end_date: str) -> pd.DataFrame:
        """获取日线数据"""
        query = """
            SELECT * FROM daily_quotes 
            WHERE symbol = :symbol AND trade_date BETWEEN :start_date AND :end_date
            ORDER BY trade_date
        """
        params = {'symbol': symbol, 'start_date': start_date, 'end_date': end_date}
        return self.db_manager.execute_query(query, params)

    def get_stock_list(self, market: str = None) -> pd.DataFrame:
        """获取股票列表"""
        if market:
            query = "SELECT * FROM stock_basic WHERE market = :market AND status = '1'"
            params = {'market': market}
        else:
            query = "SELECT * FROM stock_basic WHERE status = '1'"
            params = {}

        return self.db_manager.execute_query(query, params)

    def get_financial_data(self, symbol: str = None, year: int = None) -> pd.DataFrame:
        """获取财务数据"""
        conditions = []
        params = {}

        if symbol:
            conditions.append("symbol = :symbol")
            params['symbol'] = symbol
        if year:
            conditions.append("EXTRACT(YEAR FROM report_date) = :year")
            params['year'] = year

        where_clause = " AND ".join(conditions) if conditions else "1=1"
        query = f"""
            SELECT * FROM financial_data 
            WHERE {where_clause}
            ORDER BY report_date DESC
        """

        return self.db_manager.execute_query(query, params)

    def get_dividend_data(self, symbol: str = None, year: int = None) -> pd.DataFrame:
        """获取分红数据"""
        conditions = []
        params = {}

        if symbol:
            conditions.append("symbol = :symbol")
            params['symbol'] = symbol
        if year:
            conditions.append("EXTRACT(YEAR FROM dividend_date) = :year")
            params['year'] = year

        where_clause = " AND ".join(conditions) if conditions else "1=1"
        query = f"""
            SELECT * FROM dividend_data 
            WHERE {where_clause}
            ORDER BY dividend_date DESC
        """

        return self.db_manager.execute_query(query, params)

    def _log_update(self, data_type: str, symbol: str, records: int, status: str, error_msg: str = None):
        """记录更新日志"""
        query = """
            INSERT INTO update_log 
            (data_type, symbol, records_updated, status, error_message)
            VALUES (:data_type, :symbol, :records, :status, :error_msg)
        """
        params = {
            'data_type': data_type,
            'symbol': symbol,
            'records': records,
            'status': status,
            'error_msg': error_msg
        }
        self.db_manager.execute_query(query, params)